from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
from typing import List
import logging
import uvicorn
import numpy as np
from PIL import Image
from io import BytesIO
from magic_pdf.model.re_table.get_table_cell import ModelRunner  # 假设 ModelRunner 在 model_runner.py 文件中
from log_config import setup_logging
from main import main

# 日志初始化
setup_logging()
logger = logging.getLogger(__name__)

# FastAPI 实例
app = FastAPI()


# 初始化模型加载器（仅加载一次）
model_runner = ModelRunner()

# 请求体定义
class PDFPathsRequest(BaseModel):
    user_id: str = "test_user"
    pdf_paths: List[str]

# 预测表格标签的 API
@app.post("/predict_label/")
async def predict_label(img):
    """
    接收图像文件并返回表格标签（'wired_table' 或 'wireless_table'）。
    """

    # 获取表格标签
    label = model_runner.get_table_label(img)
    
    return {"label": label}

from config import config
LABEL_NAME = config.label_name
THRESHOLD = config.threshold
# 预测单元格坐标的 API
@app.post("/predict_cells/")
async def predict_cells(img, label_name: str = LABEL_NAME, threshold: float = THRESHOLD):
    """
    接收图像文件并返回表格单元格的坐标。
    :param label_name: 表格标签 ('wired_table' 或 'wireless_table' 或 'auto')
    :param threshold: 用于预测时过滤的置信度阈值，默认为0.3
    """
    # 获取单元格坐标
    cell_coordinates = model_runner.get_cell_coordinates(img, label_name, threshold)
    
    return {"cell_coordinates": cell_coordinates}


# API 路由
@app.post("/pdf2md/")
def pdf2md(request: PDFPathsRequest):
    logger.info(f"收到请求，user_id: {request.user_id}, pdf_paths: {request.pdf_paths}")
    try:
        md_paths = main(request.pdf_paths)
        logger.info(f"处理完成，markdown_paths: {md_paths}")
        return {"markdown_paths": md_paths}
    except Exception as e:
        logger.error(f"处理失败: {e}", exc_info=True)
        return {"error": str(e)}

if __name__ == "__main__":
    import uvicorn
    logger.info("启动 MinerU 服务...")
    from config import config
    HOST = config.host
    PORT = config.port
    uvicorn.run("api_service:app", host=HOST, port=PORT, reload=True)
